# -*- coding: utf-8 -*-
import sklearn
from numpy import genfromtxt
from numpy import transpose
import numpy
from sklearn.datasets import make_classification
from sklearn.cluster import Birch
from matplotlib import pyplot
from sklearn.cluster import DBSCAN
from numpy import where
from numpy import unique
import cv2
from PIL import ImageFont, ImageDraw, Image

rgb = [(0, 255, 0), (100, 100, 100), (255, 255, 255), (255, 0, 0), (0, 0, 255)]


def readline_count(file_name):
    return len(open(file_name).readlines())


def a():
    d0 = genfromtxt('DJI_0806.txt', delimiter=' ')
    d = transpose([d0[:, 1], d0[:, 2]])
    dbscan = DBSCAN(eps=0.02, min_samples=20).fit(d)
    yhat = dbscan.labels_
    clusters = unique(yhat)
    data = numpy.genfromtxt("DJI_0806.txt", dtype=[int, float, float, float, float])
    img = cv2.imread("DJI_0806.JPG")
    img_width = img.shape[1]
    img_height = img.shape[0]
    frame = img
    count = 0
    txt_len = readline_count("DJI_0806.txt")

    for cluster in clusters:
        rgb_index = cluster + 1
        print(cluster)
        print(rgb[rgb_index])
        row_ix = where(yhat == cluster)
        for row_ix_child in row_ix:
            print(row_ix_child)
            for xy_index in row_ix_child:
                x_centre = d[xy_index, 0]
                y_centre = d[xy_index, 1]
                print(y_centre)
                i = 0
                while i < 5:
                    frame = cv2.rectangle(frame,
                                          (int(x_centre * img_width) + i, int(y_centre * img_height) + i),
                                          (int(x_centre * img_width) + i, int(y_centre * img_height) + i),
                                          rgb[rgb_index], 3)
                    frame = cv2.rectangle(frame,
                                          (int(x_centre * img_width) - i, int(y_centre * img_height) - i),
                                          (int(x_centre * img_width) - i, int(y_centre * img_height) - i),
                                          rgb[rgb_index], 3)
                    frame = cv2.rectangle(frame,
                                          (int(x_centre * img_width) + i, int(y_centre * img_height) - i),
                                          (int(x_centre * img_width) + i, int(y_centre * img_height) - i),
                                          rgb[rgb_index], 3)
                    frame = cv2.rectangle(frame,
                                          (int(x_centre * img_width) - i, int(y_centre * img_height) + i),
                                          (int(x_centre * img_width) - i, int(y_centre * img_height) + i),
                                          rgb[rgb_index], 3)
                    i += 1

    fontpath = "font/simsun.ttc"
    font = ImageFont.truetype(fontpath, 128)
    img_pil = Image.fromarray(frame)
    draw = ImageDraw.Draw(img_pil)
    # 绘制文字信息
    draw.text((100, 350), "识别铝棒数量：" + str(count), font=font, fill=(255, 0, 0))
    bk_img = numpy.array(img_pil)
    cv2.imwrite("DJI_0806out.JPG", bk_img)


if __name__ == "__main__":
    a()
